#include "ShadersManager.h"

#include <cassert>
#include <D3D11.h>
#include <fstream>
#include <string>
#include <vector>

#include "main/D3DData.h"
#include "main/Globals.h"

#include "utils/DxErrorChecker.h"

ShadersManager::~ShadersManager() {
	for each (auto& elem in mInputLayouts) {
		elem.second->Release();
	}

	for each (auto& elem in mVertexShaders) {
		elem.second->Release();
	}

	for each (auto& elem in mPixelShaders) {
		elem.second->Release();
	}

	for each (auto& elem in mGeometryShaders) {
		elem.second->Release();
	}

	for each (auto& elem in mComputeShaders) {
		elem.second->Release();
	}

	for each (auto& elem in mHullShaders) {
		elem.second->Release();
	}

	for each (auto& elem in mDomainShaders) {
		elem.second->Release();
	}
}

ID3D11VertexShader* 
ShadersManager::addVertexShader(const std::wstring& shaderName, const std::wstring& filepath, const std::wstring* inputLayoutName, const D3D11_INPUT_ELEMENT_DESC* inputLayoutDesc, const uint32_t* descNumElems) {
	const bool createInputLayout = inputLayoutName != nullptr && inputLayoutDesc != nullptr && descNumElems != nullptr;

	// Names must be unique
	if (mVertexShaders.find(shaderName) != mVertexShaders.end()) {
		return nullptr;
	}

	if (createInputLayout && mInputLayouts.find(*inputLayoutName) != mInputLayouts.end()) {
		return nullptr;
	}
	
	// Store shader byte code, used to create a shader.
	std::vector<char> shaderByteCode;

	// Vertex shader
	storeShaderByteCode(filepath.c_str(), shaderByteCode);
	if (createInputLayout) {
		ID3D11InputLayout* inputLayout;
		buildVertexLayout(shaderByteCode, inputLayoutDesc, *descNumElems, inputLayout);
		mInputLayouts[*inputLayoutName] = inputLayout;
	}
	
	assert(Globals::gD3D->mDevice);
	ID3D11Device& device = *Globals::gD3D->mDevice;
	ID3D11VertexShader* shader = nullptr;
	const HRESULT result = device.CreateVertexShader(&shaderByteCode[0], shaderByteCode.size(), nullptr, &shader);
	DxErrorChecker(result);

	mVertexShaders[shaderName] = shader;

	return shader;
}

ID3D11PixelShader* 
ShadersManager::addPixelShader(const std::wstring& name, const std::wstring& filepath) {
	// Names must be unique
	if (mPixelShaders.find(name) != mPixelShaders.end()) {
		return nullptr;
	}

	// Store shader byte code, used to create a shader.
	std::vector<char> shaderByteCode;

	// Vertex shader
	storeShaderByteCode(filepath.c_str(), shaderByteCode);

	assert(Globals::gD3D->mDevice);
	ID3D11Device& device = *Globals::gD3D->mDevice;
	ID3D11PixelShader* shader = nullptr;
	const HRESULT result = device.CreatePixelShader(&shaderByteCode[0], shaderByteCode.size(), nullptr, &shader);
	DxErrorChecker(result);

	mPixelShaders[name] = shader;

	return shader;
}

ID3D11HullShader*
ShadersManager::addHullShader(const std::wstring& name, const std::wstring& filepath) {
	// Names must be unique
	if (mHullShaders.find(name) != mHullShaders.end()) {
		return nullptr;
	}

	// Store shader byte code, used to create a shader.
	std::vector<char> shaderByteCode;

	// Vertex shader
	storeShaderByteCode(filepath.c_str(), shaderByteCode);

	assert(Globals::gD3D->mDevice);
	ID3D11Device& device = *Globals::gD3D->mDevice;
	ID3D11HullShader* shader = nullptr;
	const HRESULT result = device.CreateHullShader(&shaderByteCode[0], shaderByteCode.size(), nullptr, &shader);
	DxErrorChecker(result);

	mHullShaders[name] = shader;

	return shader;
}

ID3D11DomainShader*
ShadersManager::addDomainShader(const std::wstring& name, const std::wstring& filepath) {
	// Names must be unique
	if (mDomainShaders.find(name) != mDomainShaders.end()) {
		return nullptr;
	}

	// Store shader byte code, used to create a shader.
	std::vector<char> shaderByteCode;

	// Vertex shader
	storeShaderByteCode(filepath.c_str(), shaderByteCode);

	assert(Globals::gD3D->mDevice);
	ID3D11Device& device = *Globals::gD3D->mDevice;
	ID3D11DomainShader* shader = nullptr;
	const HRESULT result = device.CreateDomainShader(&shaderByteCode[0], shaderByteCode.size(), nullptr, &shader);
	DxErrorChecker(result);

	mDomainShaders[name] = shader;

	return shader;
}

ID3D11GeometryShader*
ShadersManager::addGeometryShader(const std::wstring& name, const std::wstring& filepath) {
	// Names must be unique
	if (mGeometryShaders.find(name) != mGeometryShaders.end()) {
		return nullptr;
	}

	// Store shader byte code, used to create a shader.
	std::vector<char> shaderByteCode;

	// Vertex shader
	storeShaderByteCode(filepath.c_str(), shaderByteCode);

	assert(Globals::gD3D->mDevice);
	ID3D11Device& device = *Globals::gD3D->mDevice;
	ID3D11GeometryShader* shader = nullptr;
	const HRESULT result = device.CreateGeometryShader(&shaderByteCode[0], shaderByteCode.size(), nullptr, &shader);
	DxErrorChecker(result);

	mGeometryShaders[name] = shader;

	return shader;
}

ID3D11ComputeShader* 
ShadersManager::addComputeShader(const std::wstring& name, const std::wstring& filepath) {
	// Names must be unique
	if (mComputeShaders.find(name) != mComputeShaders.end()) {
		return nullptr;
	}

	// Store shader byte code, used to create a shader.
	std::vector<char> shaderByteCode;

	// Vertex shader
	storeShaderByteCode(filepath.c_str(), shaderByteCode);

	assert(Globals::gD3D->mDevice);
	ID3D11Device& device = *Globals::gD3D->mDevice;
	ID3D11ComputeShader* shader = nullptr;
	const HRESULT result = device.CreateComputeShader(&shaderByteCode[0], shaderByteCode.size(), nullptr, &shader);
	DxErrorChecker(result);

	mComputeShaders[name] = shader;

	return shader;
}

ID3D11InputLayout* 
ShadersManager::inputLayout(const std::wstring& name) const {
	InputLayouts::const_iterator it = mInputLayouts.find(name);
	assert(it != mInputLayouts.end());

	return it->second;
}

ID3D11VertexShader*
ShadersManager::vertexShader(const std::wstring& name) const {
	VertexShaders::const_iterator it = mVertexShaders.find(name);
	assert(it != mVertexShaders.end());

	return it->second;
}

ID3D11PixelShader*
ShadersManager::pixelShader(const std::wstring& name) const {
	PixelShaders::const_iterator it = mPixelShaders.find(name);
	assert(it != mPixelShaders.end());

	return it->second;
}

ID3D11HullShader*
ShadersManager::hullShader(const std::wstring& name) const {
	HullShaders::const_iterator it = mHullShaders.find(name);
	assert(it != mHullShaders.end());

	return it->second;
}

ID3D11DomainShader*
ShadersManager::domainShader(const std::wstring& name) const {
	DomainShaders::const_iterator it = mDomainShaders.find(name);
	assert(it != mDomainShaders.end());

	return it->second;
}

ID3D11GeometryShader*
ShadersManager::geometryShader(const std::wstring& name) const {
	GeometryShaders::const_iterator it = mGeometryShaders.find(name);
	assert(it != mGeometryShaders.end());

	return it->second;
}

ID3D11ComputeShader* 
ShadersManager::computeShader(const std::wstring& name) const {
	ComputeShaders::const_iterator it = mComputeShaders.find(name);
	assert(it != mComputeShaders.end());

	return it->second;
}

void 
ShadersManager::storeShaderByteCode(const std::wstring& fileName, std::vector<char>& buffer) const {
	assert(!fileName.empty());

	// Open compiled shader file and 
	// store content in the vector
	std::ifstream file;
	file.open(fileName.c_str(), std::ios::binary);
	assert(file.is_open());

	file.seekg(0, std::ios_base::end);
	size_t size = static_cast<size_t> (file.tellg());
	file.seekg(0, std::ios_base::beg);
	buffer.resize(size);
	file.read(&buffer[0], size);

	file.close();
}

void 
ShadersManager::buildVertexLayout(std::vector<char>& shaderByteCode, const D3D11_INPUT_ELEMENT_DESC* inputLayoutDesc, const uint32_t inputLayoutDescSize, ID3D11InputLayout* &inputLayout) const {
	assert(!shaderByteCode.empty());

	assert(Globals::gD3D->mDevice);
	ID3D11Device& device = *Globals::gD3D->mDevice;

	// Create the input layout        
	const HRESULT result = device.CreateInputLayout(inputLayoutDesc, inputLayoutDescSize, &shaderByteCode[0], shaderByteCode.size(), &inputLayout);
	DxErrorChecker(result);
}
